
#include "gcd_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
constexpr int32_t NUM = 12;
constexpr int VAL_ZERO = 0;
constexpr int X1_INDEX = 0;
constexpr int X2_INDEX = 1;

// broadcast params
uint32_t reduce1[20] = {VAL_ZERO};
uint32_t reduce2[20] = {VAL_ZERO};
uint32_t shape[20] = {VAL_ZERO};
uint32_t d = 1;
uint32_t dim = VAL_ZERO;

// tilling params
// uint32_t start_length = VAL_ZERO;
// uint32_t end_length = VAL_ZERO;
// uint32_t weight_length = VAL_ZERO;
// uint32_t total_length = VAL_ZERO;
// uint32_t ALIGN_NUM = VAL_ZERO;
// uint32_t tiling_size = VAL_ZERO;
// uint32_t block_size = VAL_ZERO;
// uint32_t core_size = VAL_ZERO;
// uint32_t core_remain = VAL_ZERO;
// uint32_t mode = VAL_ZERO;
    constexpr int32_t BUFFER_NUM = 2; 
    constexpr int32_t CALC_ALIGN_NUM =256;
    constexpr int32_t ALIGN_NUM = 32;  

    uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if(b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }    
    uint32_t RoundDown(uint32_t a, uint32_t b)
    { 
        if(b == 0) {
            return 0;
        }        
        return (a / b) * b;
    }  



static void GetTensorShape(const gert::StorageShape* shape, std::vector<uint64_t>& inshapeVector, uint32_t outDimNum) {
    int n = outDimNum;
    for (int j = shape->GetStorageShape().GetDimNum() - 1; j >= 0; --j) {
        inshapeVector[--n] = shape->GetStorageShape().GetDim(j);
    }
}

static void GetBroadCastParams(std::vector<uint64_t> inshapeVector1, std::vector<uint64_t> inshapeVector2,
                               gert::Shape outshape){
    // Broadcast based on the GreaterEqual operator from 0x8C
    // Caculate broadcast params
    for(int i=0;i<outshape.GetDimNum();i++){
        shape[i] = outshape.GetDim(i);
        d *= outshape.GetDim(i);
        if(inshapeVector1[i] != outshape.GetDim(i)) reduce1[i] = 1;
        if(inshapeVector2[i] != outshape.GetDim(i)) reduce2[i] = 1;
    }

    dim = outshape.GetDimNum();
    std::cout << "real_dim= " << (dim) <<std::endl;
    for(int i=dim-1;i>=1;i--){
        if(!reduce1[i - 1] && !reduce2[i - 1] && !reduce1[i] && !reduce2[i] ){
            dim--;
            shape[i - 1] *= shape[i];
        }else{
            break;
        }
    }
    if(reduce1[dim - 1] || reduce2[dim - 1]){
        shape[dim] = 1;
        dim++;
    }
}

static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    GcdTilingData tiling;

    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto aicNum = ascendcPlatform.GetCoreNumAic();
    auto aivNum = ascendcPlatform.GetCoreNumAiv();


    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

        // Get input tensors' shape
    auto outshape = context->GetOutputShape(0)->GetOriginShape();
    uint32_t outDimNum = outshape.GetDimNum();

    uint32_t totalLength = 1;
    for(int i=0;i<outDimNum;i++){
        totalLength *= outshape.GetDim(i);
    }    

    const gert::StorageShape* shape1 = context->GetInputShape(0);
    const gert::StorageShape* shape2 = context->GetInputShape(1);
    std::vector<uint64_t> inshapeVector1(outDimNum, 1);
    std::vector<uint64_t> inshapeVector2(outDimNum, 1);
    GetTensorShape(shape1, inshapeVector1, outDimNum);
    GetTensorShape(shape2, inshapeVector2, outDimNum);

    // Check if broadcasting is needed
    bool flag = false;
    for(int i=0;i<outshape.GetDimNum();i++){
        if(shape1->GetStorageShape().GetDim(i) != shape2->GetStorageShape().GetDim(i)) 
            flag = true;
    }

    auto dt = context->GetInputTensor(0)->GetDataType();
    uint32_t dataWidth;
    uint32_t ubfactor;  
    if ((dt == ge::DT_INT16)) { 
            dataWidth = 2;
            ubfactor = 3; 
    }
    else if ((dt == ge::DT_INT32)) {         
        dataWidth = 4;
        ubfactor = 3; 
    }
    else if (dt == ge::DT_INT64) {         
        dataWidth = 8;
        ubfactor = 3; 
    }     
    else {
        std::cout << "The input data type is not supported!!!" << std::endl;
        return ge::GRAPH_FAILED;
    }

    int elementsPerBlock = ALIGN_NUM / dataWidth;       
    int elementsPerRepeat = CALC_ALIGN_NUM / dataWidth;    

    int32_t totalLengthBlockAlign = RoundUp(totalLength, elementsPerBlock);
    
    aivNum = (aivNum <  totalLengthBlockAlign / elementsPerBlock) ? aivNum : totalLengthBlockAlign / elementsPerBlock;
    aivNum = (aivNum >= 1) ? aivNum : 1;

    uint32_t smallBlockLength = (totalLengthBlockAlign / aivNum / elementsPerBlock) * elementsPerBlock;
    
    
    uint32_t bigDataCoreNum = (totalLengthBlockAlign / elementsPerBlock) % aivNum;
    uint32_t bigBlockLength = 0;
    if(bigDataCoreNum > 0)
    {
        bigBlockLength = smallBlockLength + elementsPerBlock;
    }

    uint32_t maxTileCalcBlock = ub_size / (CALC_ALIGN_NUM * ubfactor) / BUFFER_NUM;
    uint32_t maxTileLength =  maxTileCalcBlock * (CALC_ALIGN_NUM / dataWidth);
    
     // uint32_t maxTileLength =64;
  
    //bigcore
    uint32_t bigTileLength = 0;
    uint32_t bigBlockLengthCalcAlign = 0;
    uint32_t bigTileNum = 0;
    uint32_t bigLasttileLength = 0;
    if(bigDataCoreNum > 0)
    {
        bigBlockLengthCalcAlign = RoundUp(bigBlockLength, elementsPerRepeat);

        (bigBlockLengthCalcAlign ) > maxTileLength ? bigTileLength = maxTileLength : bigTileLength = (bigBlockLengthCalcAlign);
        (bigTileLength >= (CALC_ALIGN_NUM * BUFFER_NUM / dataWidth)) ? bigTileLength = bigTileLength / BUFFER_NUM : bigTileLength = bigTileLength;

        bigTileNum = bigBlockLengthCalcAlign / bigTileLength;
        if (bigBlockLengthCalcAlign % bigTileLength != 0) {
            bigTileNum += 1;
        }
        bigLasttileLength = bigBlockLength - (bigTileNum -1) * bigTileLength;
    }

    //smallcore
    uint32_t smallTileLength = 0;
    uint32_t smallBlockLengthCalcAlign = 0; 
    uint32_t smallTileNum = 0;
    uint32_t smallLasttileLength = 0;

    if(bigDataCoreNum < aivNum)
    {
        smallBlockLengthCalcAlign = RoundUp(smallBlockLength, elementsPerRepeat);

        (smallBlockLengthCalcAlign ) > maxTileLength ? smallTileLength = maxTileLength : smallTileLength = (smallBlockLengthCalcAlign);
        (smallTileLength >= (CALC_ALIGN_NUM * BUFFER_NUM / dataWidth)) ? smallTileLength = smallTileLength / BUFFER_NUM : smallTileLength = smallTileLength;
        
        smallTileNum = smallBlockLengthCalcAlign / smallTileLength;
        if (smallBlockLengthCalcAlign % smallTileLength != 0) {
            smallTileNum += 1;
        }
        smallLasttileLength = smallBlockLength - (smallTileNum -1) * smallTileLength;        
    }

    if(flag) 
    {//need broadcast
        //以最后一维作为单位切分
        std::cout << "need boardcast" <<std::endl;
        context->SetTilingKey(2);
        GetBroadCastParams(inshapeVector1, inshapeVector2, outshape);
        tiling.set_totalLength(totalLength);
        tiling.set_shape(shape);
        tiling.set_reduce1(reduce1);
        tiling.set_reduce2(reduce2);
        tiling.set_dim(dim);
        context->SetBlockDim(1);

        std::cout << "totalLength= " << (totalLength) <<std::endl;
        std::cout << "dim= " << (dim) <<std::endl;

        for(int j=0; j< dim; j++) {
            std::cout << "shape[" << j << "]"<< (shape[j]) <<std::endl;
            std::cout << "reduce1[" << j << "]"<< (reduce1[j]) <<std::endl;
            std::cout << "reduce2[" << j << "]"<< (reduce2[j]) <<std::endl;
        }
    }
    else
    {//no boardcast
        std::cout << "no boardcast" <<std::endl;
        context->SetTilingKey(1);
        context->SetBlockDim(aivNum);
    }

    
    tiling.set_bigDataCoreNum(bigDataCoreNum);
    tiling.set_smallBlockLength(smallBlockLength);
    tiling.set_bigBlockLength(bigBlockLength);
    tiling.set_smallTileNum(smallTileNum);
    tiling.set_smallTileLength(smallTileLength);
    tiling.set_smallLasttileLength(smallLasttileLength);
    tiling.set_bigTileNum(bigTileNum);
    tiling.set_bigTileLength(bigTileLength);
    tiling.set_bigLasttileLength(bigLasttileLength);

    std::cout << "aivNum= " << (aivNum) <<std::endl;
    std::cout << "bigDataCoreNum= " << (bigDataCoreNum) <<std::endl;
    std::cout << "smallBlockLength= " << (smallBlockLength) <<std::endl;
    std::cout << "bigBlockLength= " << (bigBlockLength) <<std::endl;
    std::cout << "smallTileNum= " << (smallTileNum) <<std::endl;
    std::cout << "smallTileLength= " << (smallTileLength) <<std::endl;
    std::cout << "smallLasttileLength= " << (smallLasttileLength) <<std::endl;
    std::cout << "bigTileNum= " << (bigTileNum) <<std::endl;
    std::cout << "bigTileLength= " << (bigTileLength) <<std::endl;
    std::cout << "bigLasttileLength= " << (bigLasttileLength) <<std::endl;

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Gcd : public OpDef {
public:
    explicit Gcd(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(Gcd);
}
